from ultralytics import YOLO
import cv2
import glob
import os

# weight_path = r"D:\code_work\git\model_library\weight\yolo_accident_obb.pt"
# weight_path = r"E:\zhihao\2025文档\25年9月\train\weights\best.pt"
weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\服务器上的权重\nanshan-fire-v3.pt"
# weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\服务器上的权重\nanshan_fire_20251009.pt"
# 事故
# weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\public_security\traffic_accient\20250925\train77\weights\best.pt"
# video_path = r"E:\项目\松山湖公安分局无人机自动巡检项目\事故检测\无人机事故视频-云遥\0702\01.mp4"
# image_path = [r"E:\zhihao\2025文档\25年9月\红外报告\train_red0918\a4daae3413dd158e3c37f7719e396a18.png"
# image_path = [r"E:\zhihao\2025文档\25年10月\图片\1.png"]
# image_path = [r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\1.jpg",
#               r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\2.jpg",
#               r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\3.jpg",
#               r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\4.jpg",
#               r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\5.jpg",
#               r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\6.jpg",
#               r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images\7.jpg",
#               ]
# image_path = [r"E:\zhihao\2025文档\25年10月\图片\2.png"]

# 基础路径
base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images"
# 自动获取所有jpg图片
image_path = glob.glob(os.path.join(base_path, "*.jpg"))
# 如果需要按数字顺序排序
image_path.sort(key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
print(f"找到 {len(image_path)} 张图片:")
for img in image_path:
    print(img)
# 打开视频文件


model = YOLO(weight_path)
# a = model.predict(image_path, show=True, classes=None, conf=0.3, save=True, half=False, imgsz=(1920, 1080), iou=0.3)
a = model.predict(image_path, show=True, classes=None, conf=0.65, save=True, half=False, imgsz=(1920, 1080), iou=0.3)
# print(a[0].boxes.xyxy.tolist())
